import matplotlib.pyplot as plt
import random

# Create the samples
X = list(range(40))
y = []
for i in range(len(X)):
    y.append(20+X[i]+random.random()*20)

# linear regression evolution 
fig, axs = plt.subplots(1,4, figsize=(15,4))

# return the prediction of x data, by using k neighbours
def KNN_prediction(k, X, y, x, figure):
    distance = []
    for i in range(len(X)):  
        distance.append((X[i],y[i],(x-X[i])**2))  # create pairs (X[i],
                        y[i], distance from x to X[i]) 
  # sorts distance vector attending to the set of distances from x to X[i]
    sorted_list = sort(distance)  
    
    x_neighbourhood, y_neighbourhood = [], []; 
    y_prediction = 0.0
    for i in range(k):  # k-neighbours
        y_prediction += sorted_list[i][1]
        x_neighbourhood.append(sorted_list[i][0])
        y_neighbourhood.append(sorted_list[i][1])
    y_prediction /= k    

    plot(figure,X,y,'y.', "samples")    # plot samples
    # plot the k neighbours     
    plot(figure,x_neighbourhood,y_neighbourhood,'k.', "neighbourhood") 
    # plot the (x,y_prediction) prediction   
    plot(figure,x,y_prediction,"k+", 'prediction')
    axs[figure].set_xlabel('{:1.0f}'.format(k) + ' neighbors')
    return y_prediction

def sort(unsorted_list):
    return (sorted(unsorted_list, key = lambda x: x[2]))

def plot(fig, X, y, parameters, label):
    axs[fig].plot(X, y, parameters, label = label)
    axs[fig].legend(); axs[fig].grid(); 
    return

KNN_prediction(2,X, y, 7.6, 0)
KNN_prediction(4,X, y, 20.5, 1)
KNN_prediction(7,X, y, 36.2, 2)
KNN_prediction(15,X, y, 24.3, 3)

plt.show()
